import cv2
import gradio as gr
import tqdm

from utils import *
import qrcode
import io
from PIL import Image
from pyzbar import pyzbar
import time
import shutil
import os
import io
try:
    from moviepy import VideoFileClip, AudioFileClip
except:
    from moviepy.editor import VideoFileClip, AudioFileClip
import imageio

if os.name == 'nt':
    os.system('rmdir /s /q %temp%\\gradio\\')  # 删除Windows系统下的gradio缓存目录
else:
    os.system('rm -rf /tmp/gradio')  # 删除linux系统下的gradio缓存

processorMethod = ''
maxWidth = 0
def imageProcessor(img):  # 单张图片处理器
    global processorMethod, maxWidth
    if processorMethod == '自适应色阶去雾':
        img = ACED(img)
    elif processorMethod == '自适应色阶图像增强':
        img = ACE(img, 1, 10, 10)
    elif processorMethod == '直方图均衡':
        img = ehist(img)
    if maxWidth !=0:
        w,h,d = img.shape
        h = int(h*maxWidth/w)
        img = cv2.resize(img, (maxWidth, h), cv2.INTER_CUBIC)
    return img
def videoProcessor(getFrame, t):  # 视频单帧处理器
    frame = getFrame(t)
    return imageProcessor(frame)
def singleFile(Input, Method, Repeat, saveopt, quality, maxwidth, textinput):  # single file标签的处理回调函数
    global processorMethod, maxWidth
    processorMethod = Method  # 根据提交的method设置全局method
    maxWidth = maxwidth
    
    # 获取文件后缀名 using more robust method
    if isinstance(Input, str):
        basename = os.path.basename(Input)
        suffix = os.path.splitext(basename)[1][1:].lower() if '.' in basename else ''
    else:
        # Handle case where Input might be a file-like object or None
        suffix = ''
    
    # Process based on file type
    if suffix in ['mp4', 'gif']:
        video = VideoFileClip(Input)
        img = Image.fromarray(video.get_frame(0))
        video = video.fl(videoProcessor)

    else:
        img = Image.open(Input).convert('RGB')
    img = np.array(img)  # 预览文件

    for i in range(Repeat):
        img = imageProcessor(img)  # 根据全局method对单张图片进行处理
    if len(textinput) > 0:
        img = add_words(img, textinput)  # 添加表情包文字
    newPath = genNewPath(Input)  # 生成新的文件输出路径
    # np.savetxt( "a.csv", img[:,:,0], delimiter=",")  # 保存图像为csv
    if suffix in ['mp4', 'gif']:
        if saveopt == 'gif':
            newPath = ''.join(newPath.split('.')[:-1])+'.gif'
            # video.write_gif(newPath, fps=10)
            frames = []
            video.set_fps(10)
            gifPbar = tqdm.tqdm(total=int(video.duration*video.fps))
            for i in video.iter_frames():
                frames.append(i)
                gifPbar.update(1)
            imageio.mimsave(newPath, fps=video.fps, ims=frames)
        else:
            video.write_videofile(newPath, bitrate=str(quality*100)+'k')
    else:
        newPath = ''.join(newPath.split('.')[:-1])
        if saveopt == 'png':
            newPath += '.png'
            Image.fromarray(img).save(newPath)
        elif saveopt == 'jpg':
            newPath += '.jpg'
            Image.fromarray(img).save(newPath, optimize=True, quality=quality)
        else:
            newPath += '.'+suffix
            Image.fromarray(img).save(newPath)
    return newPath, img
def Swap(Output):
    return Output
def CompressCheckChange(check):
    Quality.__setattr__('interactive', check)
def InputUpload(Input):  # 输入upload回调函数
    if Input == None:
        return
def QRFunction(Image, Text, Mode):
    if Mode == '识别':
        Text = ";\r\n".join([i.data.decode('utf-8') for i in pyzbar.decode(Image)])
    elif Mode == '生成':
        Image = np.array(qrcode.make(Text).get_image().convert("L"))
        x,y = (Image - 255).nonzero()
        Image = Image[min(x):max(x), min(y):max(y)]
    return Text, Image


with gr.Blocks() as demo:
    with gr.Tab('single-file'):
        Input = gr.components.File(label='Input', file_types=['image', '.mp4', '.gif'])
        SwitchBtn = gr.Button(value='交换输入输出')
        # Method = gr.Dropdown(['直方图均衡', '自适应色阶去雾', '自适应色阶图像增强', '压缩画质'], value='自适应色阶去雾', multiselect=False)
        Method = gr.Radio(label="Method", choices=['直方图均衡', '自适应色阶去雾', '自适应色阶图像增强', '压缩画质'], value='自适应色阶去雾')
        SaveOpt = gr.Radio(label='保存文件格式', choices=['jpg', 'png', 'mp4', 'gif'], value='jpg')
        Repeat = gr.Slider(label='Repeat', value=1, minimum=1, maximum=5, step=1)
        Quality = gr.Slider(label='压缩质量', value=100, minimum=0, maximum=100, step=1)
        MaxWidth = gr.Number(label='设置最大宽度', value=0, minimum=0, step=1)
        TextInput = gr.Text(label="表情包添加文字")
        Output = gr.components.File(label='Output')
        OutputImg = gr.Image(label='Output')

        gr.Interface(fn=singleFile, inputs=[Input, Method, Repeat, SaveOpt, Quality, MaxWidth, TextInput], outputs=[Output, OutputImg])
        SwitchBtn.click(fn=Swap, inputs=[Output], outputs=[Input])
        Input.change(fn=InputUpload, inputs=[Input])

    with gr.Tab("QR Code"):
        QRTexti = gr.Text(label="QRTextIn")
        QRImagei = gr.Image(label='ImageIn')
        QRImageo = gr.Image(label='ImageOut',format='jpeg', image_mode='L')
        QRTexto = gr.Text(label='QRTextOut')
        QRMode = gr.Radio(label='mode', choices=['识别', '生成'], value='识别')
        gr.Interface(fn=QRFunction,
                     inputs=[QRImagei, QRTexti, QRMode], outputs=[QRTexto, QRImageo])
    with gr.Tab('PDF'):
        pdfFile = gr.File(label='PDF File', file_types=['.pdf'])
        zipOut = gr.File(label='Output')
        Preview = gr.Gallery()
        x = gr.Slider(label='放缩', value=3.5, minimum=0.2, maximum=10, step=0.5)
        gr.Interface(fn=pdf2img, inputs=[pdfFile, x], outputs=[zipOut, Preview])
    with gr.Tab("RegionGrowing"):
        ckbox = gr.Checkbox(label='直方图均衡预处理', value=False)
        ckbox2 = gr.Checkbox(label='canny边缘提取', value=False)
        erode_iter = gr.Slider(label='mask腐蚀迭代步数', value=3, maximum=8, step=1)
        x = gr.Slider(label='放缩', value=0.25, minimum=0.1, maximum=1, step=0.05)
        threshold = gr.Slider(label='threshold', value=1, minimum=1, maximum=80, step=1)
        with gr.Row():
            with gr.Column():
                im = gr.ImageEditor(
                type="numpy",
                label='添加mask'
            )
                btn1 = gr.Button(value='运行区域生长')
                btn2 = gr.Button(value='叠加mask')

            mask = gr.Image(visible=False)
            im_preview = gr.Image(format='png', label='preview', image_mode='RGBA')
            region_res = gr.Image(label='region_mask', interactive=False)
            org_im = gr.Image(visible=False)
        im.upload(upload, outputs=[org_im, im_preview], inputs=im, show_progress="hidden")
        im.change(get_mask, outputs=mask, inputs=im, show_progress="hidden")
        btn1.click(fn=run_region_grow, inputs=[im, org_im, ckbox, ckbox2, erode_iter, x, threshold], outputs=[region_res, im_preview])
        btn2.click(fn=swap, inputs=[org_im,region_res], outputs=im)
    with gr.Tab("GetAudio"):
        videoFile = gr.File(label='mp4 File', file_types=['.mp4'])
        audioFile = gr.File(label='Audio File')
        gr.Interface(fn=getAudio, inputs=[videoFile], outputs=[audioFile])


demo.launch(server_name='0.0.0.0')
